import matplotlib.pyplot as plt
import numpy as np
from datetime import datetime

class TrainingVisualizer:
    """
    训练过程可视化器
    支持训练损失、准确率和学习率的实时可视化
    """
    
    def __init__(self, parent_frame=None, use_tkagg=False, figsize=(8, 6), plot_type='both'):
        """
        初始化训练可视化器
        
        参数:
        - parent_frame: Tkinter父框架
        - use_tkagg: 是否使用TkAgg后端
        - figsize: 图表大小
        - plot_type: 图表类型，可选值: 'loss', 'accuracy', 'both'
        """
        self.parent_frame = parent_frame
        self.use_tkagg = use_tkagg
        self.figsize = figsize
        self.plot_type = plot_type  # 'loss', 'accuracy', 'both'
        
        # 初始化数据
        self.epochs = []
        self.train_loss = []
        self.test_loss = []
        self.train_acc = []
        self.test_acc = []
        
        # 如果使用TkAgg后端，确保正确设置
        if self.use_tkagg:
            import matplotlib
            matplotlib.use('TkAgg')
            from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
            self.FigureCanvasTkAgg = FigureCanvasTkAgg
        
        # 创建图表
        self._create_figure()
    
    def _create_figure(self):
        """
        创建图表和子图
        """
        # 根据plot_type决定创建多少个子图
        if self.plot_type == 'both':
            # 创建图表 - 包含loss和accuracy两个子图
            self.fig, self.axes = plt.subplots(2, 1, figsize=self.figsize, dpi=100, sharex=True)
            self.axes = self.axes.flatten()
            
            # 初始化线条
            self.train_loss_line, = self.axes[0].plot([], [], 'b-', label='训练损失')
            self.test_loss_line, = self.axes[0].plot([], [], 'r-', label='测试损失')
            self.train_acc_line, = self.axes[1].plot([], [], 'b-', label='训练准确率')
            self.test_acc_line, = self.axes[1].plot([], [], 'r-', label='测试准确率')
            
            # 配置图表
            self.axes[0].set_ylabel('损失值')
            self.axes[0].legend()
            self.axes[0].grid(True, alpha=0.3)
            
            self.axes[1].set_xlabel('轮次')
            self.axes[1].set_ylabel('准确率 (%)')
            self.axes[1].set_ylim(0, 105)  # 准确率范围设置为0-105%
            self.axes[1].legend()
            self.axes[1].grid(True, alpha=0.3)
        elif self.plot_type == 'loss':
            # 只创建损失曲线
            self.fig, self.ax = plt.subplots(1, 1, figsize=self.figsize, dpi=100)
            self.axes = [self.ax]
            
            # 初始化线条
            self.train_loss_line, = self.ax.plot([], [], 'b-', label='训练损失')
            self.test_loss_line, = self.ax.plot([], [], 'r-', label='测试损失')
            
            # 配置图表
            self.ax.set_xlabel('轮次')
            self.ax.set_ylabel('损失值')
            self.ax.legend()
            self.ax.grid(True, alpha=0.3)
        elif self.plot_type == 'accuracy':
            # 只创建准确率曲线
            self.fig, self.ax = plt.subplots(1, 1, figsize=self.figsize, dpi=100)
            self.axes = [self.ax]
            
            # 初始化线条
            self.train_acc_line, = self.ax.plot([], [], 'b-', label='训练准确率')
            self.test_acc_line, = self.ax.plot([], [], 'r-', label='测试准确率')
            
            # 配置图表
            self.ax.set_xlabel('轮次')
            self.ax.set_ylabel('准确率 (%)')
            self.ax.set_ylim(0, 105)  # 准确率范围设置为0-105%
            self.ax.legend()
            self.ax.grid(True, alpha=0.3)
        
        # 设置图表标题
        title = 'MNIST手写数字识别 - '  
        if self.plot_type == 'both':
            title += '训练过程可视化'
        elif self.plot_type == 'loss':
            title += '损失曲线'
        elif self.plot_type == 'accuracy':
            title += '准确率曲线'
        self.fig.suptitle(title, fontsize=14)
        
        # 如果使用Tkinter嵌入
        if self.use_tkagg and self.parent_frame:
            self.canvas = self.FigureCanvasTkAgg(self.fig, master=self.parent_frame)
            self.canvas.draw()
            self.canvas.get_tk_widget().pack(fill='both', expand=True)
        
        # 调整布局
        self.fig.tight_layout()
        self.fig.subplots_adjust(top=0.92)
    
    def update(self, history_data):
        """
        更新可视化数据
        
        参数:
        - history_data: 包含训练历史的字典，应该包含以下键：
          'epoch', 'train_loss', 'test_loss', 'train_acc', 'test_acc'
        """
        # 更新数据
        self.epochs.append(history_data['epoch'])
        self.train_loss.append(history_data['train_loss'])
        self.test_loss.append(history_data['test_loss'])
        self.train_acc.append(history_data['train_acc'])
        self.test_acc.append(history_data['test_acc'])
        
        # 根据plot_type更新相应的图表
        if self.plot_type == 'loss' or self.plot_type == 'both':
            self.train_loss_line.set_data(self.epochs, self.train_loss)
            self.test_loss_line.set_data(self.epochs, self.test_loss)
        
        if self.plot_type == 'accuracy' or self.plot_type == 'both':
            self.train_acc_line.set_data(self.epochs, self.train_acc)
            self.test_acc_line.set_data(self.epochs, self.test_acc)
        
        # 更新坐标轴范围
        for ax in self.axes:
            ax.relim()
            ax.autoscale_view()
        
        # 更新图表
        self.fig.canvas.draw()
        self.fig.canvas.flush_events()
    
    def save_plot(self, filename=None):
        """
        保存图表到文件
        
        参数:
        - filename: 文件名，如果为None则自动生成
        
        返回:
        - 保存的文件路径
        """
        if filename is None:
            # 生成带有时间戳的文件名
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"training_visualization_{timestamp}.png"
        
        # 保存图表
        self.fig.savefig(filename, dpi=300, bbox_inches='tight')
        return filename
    
    def clear(self):
        """
        清除图表数据
        """
        # 重置数据
        self.epochs = []
        self.train_loss = []
        self.test_loss = []
        self.train_acc = []
        self.test_acc = []
        
        # 根据plot_type重置相应的线条
        if self.plot_type == 'loss' or self.plot_type == 'both':
            self.train_loss_line.set_data([], [])
            self.test_loss_line.set_data([], [])
        
        if self.plot_type == 'accuracy' or self.plot_type == 'both':
            self.train_acc_line.set_data([], [])
            self.test_acc_line.set_data([], [])
        
        # 更新坐标轴
        for ax in self.axes:
            ax.relim()
            ax.autoscale_view()
        
        # 更新图表
        self.fig.canvas.draw()
        self.fig.canvas.flush_events()

class ActivationVisualizer:
    """
    激活值可视化器
    支持卷积层激活值可视化和全连接层激活值可视化
    """
    
    def __init__(self, parent_frame=None, max_channels=6, cmap='viridis', show_colorbar=True, use_tkagg=True):
        """
        初始化激活值可视化器
        
        参数:
        - parent_frame: Tkinter父框架
        - max_channels: 最多显示的通道数
        - cmap: 颜色映射
        - show_colorbar: 是否显示颜色条
        - use_tkagg: 是否使用TkAgg后端
        """
        self.parent_frame = parent_frame
        self.max_channels = max_channels
        self.cmap = cmap
        self.show_colorbar = show_colorbar
        self.use_tkagg = use_tkagg
        self.figsize = (10, 8)
        
        # 如果使用TkAgg后端，确保正确设置
        if self.use_tkagg:
            import matplotlib
            matplotlib.use('TkAgg')
            from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
            self.FigureCanvasTkAgg = FigureCanvasTkAgg
        
        # 创建图表
        self._create_figure()
    
    def _create_figure(self):
        """
        创建图表
        """
        self.fig, self.ax = plt.subplots(figsize=self.figsize, dpi=100)
        
        # 如果使用Tkinter嵌入
        if self.use_tkagg and self.parent_frame:
            self.canvas = self.FigureCanvasTkAgg(self.fig, master=self.parent_frame)
            self.canvas.draw()
            self.canvas.get_tk_widget().pack(fill='both', expand=True)
    
    def visualize_convolutional_activation(self, activation_map, layer_name=None):
        """
        可视化卷积层激活值
        
        参数:
        - activation_map: 激活值映射，形状为 [channels, height, width]
        - layer_name: 层名称
        """
        # 清除当前图表
        self.ax.clear()
        self.fig.clf()
        
        # 确保输入形状正确
        if len(activation_map.shape) != 3:
            raise ValueError("激活值映射应为三维形状 [channels, height, width]")
        
        # 获取通道数
        num_channels = activation_map.shape[0]
        
        # 限制显示的通道数
        display_channels = min(num_channels, self.max_channels)
        
        # 计算网格大小
        cols = min(display_channels, 4)
        rows = (display_channels + cols - 1) // cols
        
        # 创建新的子图
        self.fig, self.axes = plt.subplots(rows, cols, figsize=self.figsize, dpi=100)
        self.axes = np.atleast_2d(self.axes)
        
        # 显示每个通道
        for i in range(display_channels):
            row = i // cols
            col = i % cols
            ax = self.axes[row, col]
            
            # 显示激活值
            im = ax.imshow(activation_map[i], cmap=self.cmap)
            ax.set_title(f"通道 {i+1}")
            ax.axis('off')
            
            # 显示颜色条
            if self.show_colorbar:
                self.fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
        
        # 隐藏未使用的子图
        for i in range(display_channels, rows * cols):
            row = i // cols
            col = i % cols
            self.axes[row, col].axis('off')
        
        # 设置图表标题
        if layer_name:
            self.fig.suptitle(f"卷积层激活值可视化 - {layer_name}", fontsize=14)
        else:
            self.fig.suptitle("卷积层激活值可视化", fontsize=14)
        
        # 调整布局
        self.fig.tight_layout()
        self.fig.subplots_adjust(top=0.92)
        
        # 如果使用Tkinter，更新canvas
        if self.use_tkagg and self.parent_frame:
            self.canvas.figure = self.fig
            self.canvas.draw()
        else:
            plt.draw()
            plt.pause(0.01)
    
    def visualize_fully_connected_activation(self, activation_values, layer_name=None):
        """
        可视化全连接层激活值
        
        参数:
        - activation_values: 激活值，形状为 [num_neurons]
        - layer_name: 层名称
        """
        # 清除当前图表
        self.ax.clear()
        
        # 确保输入形状正确
        if len(activation_values.shape) != 1:
            activation_values = activation_values.flatten()
        
        # 显示直方图
        self.ax.hist(activation_values, bins=50, alpha=0.7, color='blue', edgecolor='black')
        self.ax.set_xlabel('激活值')
        self.ax.set_ylabel('频率')
        
        # 设置图表标题
        if layer_name:
            self.ax.set_title(f"全连接层激活值分布 - {layer_name}", fontsize=12)
        else:
            self.ax.set_title("全连接层激活值分布", fontsize=12)
        
        # 添加网格
        self.ax.grid(True, alpha=0.3)
        
        # 添加统计信息
        mean_val = np.mean(activation_values)
        std_val = np.std(activation_values)
        max_val = np.max(activation_values)
        min_val = np.min(activation_values)
        
        stats_text = f"均值: {mean_val:.4f}\n标准差: {std_val:.4f}\n最大值: {max_val:.4f}\n最小值: {min_val:.4f}"
        self.ax.text(0.05, 0.95, stats_text, transform=self.ax.transAxes, 
                    verticalalignment='top', bbox=dict(boxstyle='round', facecolor='white', alpha=0.5))
        
        # 调整布局
        self.fig.tight_layout()
        
        # 如果使用Tkinter，更新canvas
        if self.use_tkagg and self.parent_frame:
            self.canvas.draw()
        else:
            plt.draw()
            plt.pause(0.01)
    
    def visualize_input_image(self, image, true_label=None, predicted_label=None):
        """
        可视化输入图像
        
        参数:
        - image: 输入图像，形状为 [height, width] 或 [3, height, width]（RGB）
        - true_label: 真实标签
        - predicted_label: 预测标签
        """
        # 清除当前图表
        self.ax.clear()
        
        # 调整图像形状
        if len(image.shape) == 3:
            # 如果是RGB图像，调整通道顺序
            if image.shape[0] == 3:
                image = np.transpose(image, (1, 2, 0))
            self.ax.imshow(image)
        else:
            # 灰度图像
            self.ax.imshow(image, cmap='gray')
        
        # 设置标题
        title_parts = []
        if true_label is not None:
            title_parts.append(f"真实标签: {true_label}")
        if predicted_label is not None:
            title_parts.append(f"预测标签: {predicted_label}")
        
        if title_parts:
            self.ax.set_title(" | ".join(title_parts), fontsize=12)
        
        # 隐藏坐标轴
        self.ax.axis('off')
        
        # 调整布局
        self.fig.tight_layout()
        
        # 如果使用Tkinter，更新canvas
        if self.use_tkagg and self.parent_frame:
            self.canvas.draw()
        else:
            plt.draw()
            plt.pause(0.01)
    
    def save_visualization(self, filename=None):
        """
        保存当前可视化到文件
        
        参数:
        - filename: 文件名，如果为None则自动生成
        
        返回:
        - 保存的文件路径
        """
        if filename is None:
            # 生成带有时间戳的文件名
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"activation_visualization_{timestamp}.png"
        
        # 保存图表
        self.fig.savefig(filename, dpi=300, bbox_inches='tight')
        return filename
    
    def clear(self):
        """
        清除当前可视化
        """
        self.ax.clear()
        
        # 如果使用Tkinter，更新canvas
        if self.use_tkagg and self.parent_frame:
            self.canvas.draw()
        else:
            plt.draw()
            plt.pause(0.01)
    
    def update(self, input_image, activations, layer_names=None):
        """
        统一的更新接口，处理不同类型的激活值可视化
        
        参数:
        - input_image: 输入图像张量
        - activations: 各层激活值列表
        - layer_names: 各层名称列表
        """
        # 先显示输入图像
        if input_image is not None:
            # 转换为numpy格式
            img_np = input_image.squeeze().cpu().detach().numpy()
            self.visualize_input_image(img_np)
        
        # 如果有激活值和层名称，依次可视化每一层
        if activations and layer_names:
            for i, (activation, layer_name) in enumerate(zip(activations, layer_names)):
                # 转换为numpy格式
                activation_np = activation.squeeze().cpu().detach().numpy()
                
                # 根据形状判断层类型
                if len(activation_np.shape) == 3:  # 卷积层 [channels, height, width]
                    self.visualize_convolutional_activation(activation_np, layer_name)
                elif len(activation_np.shape) == 1:  # 全连接层 [num_neurons]
                    self.visualize_fully_connected_activation(activation_np, layer_name)
